import numpy as np
import matplotlib.pyplot as plt
import matplotlib

class Partial:
    def __init__(self):
        self.dx = 0.18
        self.x_num = 50
        self.dt = 0.01
        self.t_num = 300
        self.arg = self.dt/(self.dx**2)
        self.U = np.zeros([self.x_num+1, self.t_num+1])
        self.space_ls = np.arange(0, (self.x_num + 1) * self.dx, self.dx)  # 建立空间等差数列，从0到3，公差是h
        self.time_ls = np.arange(0, (self.t_num + 1) * self.dt, self.dt)  # 建立空间等差数列，从0到3，公差是h

    def calculator(self):
        for t in range(self.t_num+1):
            self.U[0, t] = 20
            self.U[self.x_num, t] = 20
        # 设置初值
        for x in range(self.x_num):
            q = x*self.dx
            self.U[x, 0] = 3*q**1.5
        for t in range(self.t_num):
            for x in range(self.x_num):
                self.U[x, t + 1] = self.arg * self.U[x + 1, t] + (1 - 2 * self.arg) * self.U[x, t] + self.arg * self.U[x - 1, t]

    def draw(self):
        plt.plot(self.space_ls, self.U[:, 50], 'g-', label='t=0.5', linewidth=1.0)
        plt.plot(self.space_ls, self.U[:, 100], 'b-', label='t=1', linewidth=1.0)
        plt.plot(self.space_ls, self.U[:, 150], 'k-', label='t=1.5', linewidth=1.0)
        plt.plot(self.space_ls, self.U[:, 200], 'r-', label='t=2', linewidth=1.0)
        plt.ylabel(r'$\varphi (x,t)$', fontsize=20)
        plt.xlabel('x', fontsize=20)
        plt.xlim(0, self.x_num * self.dx)
        plt.title('位置变化曲线')
        plt.savefig('位置变化曲线.jpg')
        plt.show()

        plt.plot(self.time_ls, self.U[20, :], 'g-', label='x=%.1f'%self.dx*20, linewidth=1.0)
        plt.plot(self.time_ls, self.U[25, :], 'b-', label='x=%.1f'%self.dx*25, linewidth=1.0)
        plt.plot(self.time_ls, self.U[30, :], 'k-', label='x=%.1f'%self.dx*30, linewidth=1.0)
        plt.plot(self.time_ls, self.U[40, :], 'r-', label='x=%.1f'%self.dx*40, linewidth=1.0)
        plt.ylabel(r'$\varphi (x,t)$', fontsize=20)
        plt.xlabel('t', fontsize=20)
        plt.xlim(0, self.t_num*self.dt)
        plt.title('时间变化曲线')
        plt.savefig('时间变化曲线.jpg')
        plt.show()

        extent = [0, self.dt*self.t_num, 0, self.dx*self.x_num]  # 时间和空间的取值范围
        lv = np.arange(0, 100, 0.1)  # 温度等高线的变化范围0-10，变化间隔为0.1
        plt.contourf(self.U, lv, origin='lower', extent=extent, cmap=plt.cm.jet)
        plt.ylabel('x', fontsize=20)
        plt.xlabel('t', fontsize=20)
        plt.title('整体热成像')
        plt.savefig('整体热成像.jpg')
        plt.show()


if __name__ == '__main__':
    font = {'family': 'SimHei'}
    matplotlib.rc('font', **font)
    p = Partial()
    p.calculator()
    p.draw()